package com.example.javamaday.tree;

/**
 * 红黑树的实现
 */
public class RedBlackTree<Key extends Comparable<Key>, Value> {

    public static void main(String[] args) {
        //创建红黑树
        RedBlackTree<String, String> tree = new RedBlackTree<>();

        //往树中插入元素
        tree.put("1","张三");
        tree.put("2","李四");
        tree.put("3","王五");
        //从树中获取元素
        String r1 = tree.get("1");
        System.out.println(r1);

        String r2 = tree.get("2");
        System.out.println(r2);

        String r3 = tree.get("3");
        System.out.println(r3);
    }

    private static final boolean RED = true;  // 红色链接标识
    private static final boolean BLACK = false;  // 黑色链接标识

    private Node root;  //  记录根结点
    private int N; // 记录树中元素的个数


    // 左旋转
    private Node rotateLeft(Node h) {
        // 获取h节点的右子结点，标为x
        Node x = h.right;
        // 让x结点的左子结点为h的右子结点
        h.right = x.left;
        // 让h称为x结点的左子结点
        x.left = h;
        // 让x的color属性 = h的color属性
        x.color = h.color;
        // 让h的color = RED
        h.color = RED;
        return x;
    }

    // 右旋转
    private Node rotateRight(Node h) {
        // 获取h结点的左子结点，标为x
        Node x = h.left;
        // 让x节点的右子结点称为h结点的左子结点
        h.left = x.right;
        // 让h结点称为x结点的右子结点
        x.right = h;
        // 让x的color = h的color
        x.color = h.color;
        // 让h的color = RED
        h.color = RED;
        return x;
    }

    // 颜色反转，相当于拆分4-结点
    private void flipColors(Node h) {
        // 当前结点变为红色
        h.color = RED;
        // 左子结点和右子结点变为黑色
        h.left.color = BLACK;
        h.right.color = BLACK;
    }

    // API - 插入操作
    public void put(Key key, Value value) {
        root = put(root, key, value);
        // 保证根结点的颜色总为黑色
        root.color = BLACK;
    }

    // 内部 - 向指定树插入元素，并返回添加元素后的新树
    private Node put(Node h, Key key, Value value) {
        // 判断h是否为空，如果为空，则直接返回一个红色的结点
        if (h == null) {
            N++;
            return new Node(key, value, null, null, RED);
        }
        // 比较h的键和k的大小
        int cmp = key.compareTo(h.key);
        if (cmp < 0) {
            // 继续往左
            h.left = put(h.left, key, value);
        } else if (cmp > 0) {
            // 继续往右
            h.right = put(h.right, key, value);
        } else if (cmp == 0) {
            // 替换
            h.value = value;
        }
        // 进行左旋
        if(isRed(h.right) && !isRed(h.left)) {
            h = rotateLeft(h);
        }
        // 进行右旋
        if (isRed(h.left) && isRed(h.left.left)) {
            h = rotateRight(h);
        }
        // 进行颜色反转
        if (isRed(h.left) && isRed(h.right)) {
            flipColors(h);
        }
        return h;
    }

    // 判断当前结点的父节点指向链接是否为红色
    private boolean isRed(Node x) {
        if (x == null) {
            return false;
        }
        return x.color == RED;
    }

    // 获取树中元素的个数
    private int getSize() {
        return N;
    }

    //根据key，从树中找出对应的值
    public Value get(Key key) {
        return get(root,key);
    }

    //从指定的树x中，查找key对应的值
    public Value get(Node x, Key key) {
        if (x == null){
            return null;
        }

        //比较x结点的键和key的大小
        int cmp = key.compareTo(x.key);
        if (cmp<0){
            return get(x.left,key);
        }else if (cmp>0){
            return get(x.right,key);
        }else{
            return x.value;
        }
    }

    private class Node {

        public Key key; // 存储结点的键
        public Value value; // 存储结点的值

        public Node left; // 记录左子节点
        public Node right; // 记录右子节点

        public boolean color; // 记录其父结点指向它的链接的颜色

        public Node(Key key, Value value, Node left, Node right, boolean color) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
            this.color = color;
        }
    }

}
